package com.jzlin.funhub.admin.service.chatai;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.jzlin.funhub.framework.core.exception.AppException;
import com.jzlin.funhub.framework.core.sse.SseClient;
import com.jzlin.funhub.admin.service.chatai.model.enums.ChatAiMessageRole;
import com.jzlin.funhub.admin.service.chatai.store.ChatAiMessageRecord;
import com.jzlin.funhub.admin.service.chatai.support.AbstractChatAiService;
import com.jzlin.funhub.admin.service.chatai.support.BaseChatMessage;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.deserialize.MessageDeserializeFactory;
import com.zhipu.oapi.service.v4.model.*;
import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;

/**
 * 智谱AI大模型问答服务
 *
 * @author Lin
 * @since 2024/10/9
 */
@Service
@Slf4j
public class ZhipuChatAiService extends AbstractChatAiService {
    // 默认客户端
    private static final ClientV4 defaultClient = new ClientV4.Builder("c0ab3750647f7c7687dee058ca288594.6ynTz60hOaDCPqV0")
            .networkConfig(60, 30, 30, 30, TimeUnit.SECONDS)
            .connectionPool(new okhttp3.ConnectionPool(8, 30, TimeUnit.SECONDS))
            .build();
    private static final ObjectMapper mapper = MessageDeserializeFactory.defaultObjectMapper();

    /**
     * 流式通信
     *
     * @param message 消息
     */
    @Override
    protected void handleBySse(BaseChatMessage message) {
        try {
            List<ChatMessage> messages = buildChatMessageList(message);
            // 函数调用参数构建部分
            List<ChatTool> chatToolList = new ArrayList<>();
            ChatTool chatTool = new ChatTool();

            chatTool.setType("code_interpreter");
//            chatTool.setType("web_search");
//            WebSearch webSearch = new WebSearch();
//            webSearch.setEnable(true);
//            chatTool.setWeb_search(webSearch);
//        ObjectNode objectNode =  mapper.createObjectNode();
//        objectNode.put("code", "北京天气");
//        chatTool.set(chatFunction);
            chatToolList.add(chatTool);


             ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                    .model("glm-4-flash")
                    .stream(Boolean.TRUE)
                    .invokeMethod(Constants.invokeMethod)
                    .messages(messages)
                    .tools(chatToolList)
                    .toolChoice("auto")
                    .requestId(message.getSessionId())
                    .build();
            ModelApiResponse sseModelApiResp = defaultClient.invokeModelApi(chatCompletionRequest);
            if (sseModelApiResp.isSuccess()) {
                AtomicBoolean isFirst = new AtomicBoolean(true);
                List<Choice> choices = new ArrayList<>();
                AtomicReference<ChatMessageAccumulator> lastAccumulator = new AtomicReference<>();

                StringBuilder replyContent = new StringBuilder();
                Long startTime = System.currentTimeMillis();
                mapStreamToAccumulator(sseModelApiResp.getFlowable())
                        .doOnNext(accumulator -> {
                            {
                                if (isFirst.getAndSet(false)) {
                                    log.info("Response: ");
                                }
                                if (accumulator.getDelta() != null && accumulator.getDelta().getTool_calls() != null) {
                                    String jsonString = mapper.writeValueAsString(accumulator.getDelta().getTool_calls());
                                    log.info("tool_calls: {}", jsonString);
                                }
                                if (accumulator.getDelta() != null && accumulator.getDelta().getContent() != null) {
                                    String content = accumulator.getDelta().getContent();
                                    //log.info(content);
                                    replyContent.append(content);
                                    if (!SseClient.sendMessage(message.getSseId(), transferSseContent(content))) {
                                        throw new AppException(500, "SSE流消息推送失败");
                                    }
                                }
                                choices.add(accumulator.getChoice());
                                lastAccumulator.set(accumulator);

                            }
                        })
                        .doOnComplete(() -> {
                            log.info("Zhipu stream completed.");
                            storeChatMessage(message.getSessionId(), replyContent.toString(), ChatAiMessageRole.ASSISTANT, startTime, System.currentTimeMillis());
                            SseClient.finishSse(message.getSseId());
                        })
                        .doOnError(throwable -> {
                            log.error("Zhipu stream error", throwable);
                            SseClient.closeSse(message.getSseId());
                        }) // Handle errors
                        .blockingSubscribe();// Use blockingSubscribe instead of blockingGet()

                ChatMessageAccumulator chatMessageAccumulator = lastAccumulator.get();
                ModelData data = new ModelData();
                data.setChoices(choices);
                if (chatMessageAccumulator != null) {
                    data.setUsage(chatMessageAccumulator.getUsage());
                    data.setId(chatMessageAccumulator.getId());
                    data.setCreated(chatMessageAccumulator.getCreated());
                }
                data.setRequestId(chatCompletionRequest.getRequestId());
                sseModelApiResp.setFlowable(null);// 打印前置空
                sseModelApiResp.setData(data);
            }
            log.info("model output: {}", mapper.writeValueAsString(sseModelApiResp));
            //client.getConfig().getHttpClient().dispatcher().executorService().shutdown();

            //client.getConfig().getHttpClient().connectionPool().evictAll();
            // List all active threads
//            for (Thread t : Thread.getAllStackTraces().keySet()) {
//                log.info("Thread: " + t.getName() + " State: " + t.getState());
//            }
        } catch (Exception e) {
            log.error("问答发生异常", e);
        } finally {
            SseClient.closeSse(message.getSseId());
        }
    }

    /**
     * 构建消息
     * @return
     */
    private List<ChatMessage> buildChatMessageList(BaseChatMessage message) {
        // 历史消息记录
        List<ChatAiMessageRecord> historyMessages = getChatRefContext(message.getSessionId()).getHistoryMessage();

        return historyMessages.stream().map(this::convertChatMessage).collect(Collectors.toList());
    }

    private ChatMessage convertChatMessage(ChatAiMessageRecord record) {
        ChatMessage message = new ChatMessage();
        message.setContent(record.getContent());
        switch (ChatAiMessageRole.of(record.getRole())) {
            case USER:
                message.setRole(ChatMessageRole.USER.value());
                break;
            case ASSISTANT:
                message.setRole(ChatMessageRole.ASSISTANT.value());
                break;
            case SYSTEM:
                message.setRole(ChatMessageRole.SYSTEM.value());
                break;
            case FUNCTION:
                message.setRole(ChatMessageRole.FUNCTION.value());
                break;

        }

        return message;
    }

    private Flowable<ChatMessageAccumulator> mapStreamToAccumulator(Flowable<ModelData> flowable) {
        return flowable.map(chunk -> new ChatMessageAccumulator(chunk.getChoices().get(0).getDelta(), null, chunk.getChoices().get(0), chunk.getUsage(), chunk.getCreated(), chunk.getId()));
    }
}
